rm(list = ls())

set.seed(1)

library(minqa)
library(VGAM)
library(gtools)

setwd("~/Dropbox/complieropt/Polmeth Archive")
source("icsw.R")

###

dgp <- function(Z,X1,X2,X3,a,b,c,d,e,f,g,h) {
	N <- length(Z)
	AC <- (2 + a + b*X1 + c*X2 + d*X3 + rnorm(N,0,3) > 0)
	A <- (-2 + e + f*X1 + g*X2 + h*X3 + rnorm(N,0,3) > 0)*AC
	C <- AC - A
	
	D <- C*Z + A
	Y0 <- 5*AC + rnorm(N)
	tau<- 5*(1+X1+X2+X3)
	Y <- Y0 + tau*(D)
	
	return(list(D,Y,tau,C))
	}

perm.test <- function(D,Z,Xmat,cscoreout,numperm=50,genoud=FALSE) {
	test.dist <- rep(NA,numperm)
	
	for (i in 1:numperm) {
		Xmats <- Xmat[sample(1:length(D)),]
		cout <- complier.score(D,Z,Xmats,genoud=genoud)
		test.dist[i] <- sum((D - cout$C.score*Z - cout$A.score)^2)
	}
	
	test.stat <- sum((D - cscoreout$C.score*Z - cscoreout$A.score)^2)
	p.value <- mean(test.dist <= test.stat)
	output <- list(test.dist,test.stat,p.value)
	names(output) <- c("test.dist","test.stat","p.value")
	return(output)
}

# ATE is always 1

iteration <- function(N,alphalist) {

Z <- sample(c(rep(1,N/2),rep(0,N/2)))

X1 <- runif(N,-1,1)
X2 <- runif(N,-1,1)
X3 <- runif(N,-1,1)

a <- runif(1,-2,2)
b <- runif(1,-2,2)
c <- runif(1,-2,2)
d <- runif(1,-2,2)
e <- runif(1,-2,2)
f <- runif(1,-2,2)
g <- runif(1,-2,2)
h <- runif(1,-2,2)

output <- rep(NA,6)

# DGP1: Correct Model

dgpout <- dgp(Z,X1,X2,X3,a,b,c,d,e,f,g,h)

D <- dgpout[[1]]
Y <- dgpout[[2]]

Xmat <- W <- cbind(X1,X2,X3)

cscoreout <- complier.score(genoud=FALSE,D,Z,Xmat)
cscore <- cscoreout$C.score
cscore[cscore < .Machine$double.eps^.5] <- .Machine$double.eps^.5

output[1] <- mean(dgpout[[3]])
output[2] <- lm(Y~Z+Xmat)$coefficients["Z"]
output[3] <- lm(Y~D+Xmat)$coefficients["D"]
output[4] <- tsls.wfit(cbind(1,D,Xmat),Y,cbind(1,Z,Xmat),weights=rep(1,N))$coefficients["D"]
output[5] <- tsls.wfit(cbind(1,D,Xmat),Y,cbind(1,Z,Xmat),weights=1/cscore)$coefficients["D"]
output[6] <- mean(dgpout[[4]])

llist <- length(alphalist)
icswcoefs <- rep(NA,llist)

for (i in 1:llist) {
	alpha <- alphalist[i]
	cscoreA <- cscore
	qcscore <- quantile(cscore,1/N^alpha)
	cscoreA[cscoreA < qcscore] <- qcscore
	icswcoefs[i] <- tsls.wfit(cbind(1,D,Xmat),Y,cbind(1,Z,Xmat),weights=1/cscoreA)$coefficients["D"]
	}
	
return(c(output,icswcoefs))
}

numiter <- 2500

#N: 500, 1000, 2500, 5000, 10000, 25000

N <- 25000

alphalist <- seq(0.00,0.5,0.005)

simoutput <- sapply(rep(N,numiter),iteration,alphalist=alphalist)

save.image(paste("alpha ",N,".rdata",sep=""))
